import os
from dotenv import load_dotenv
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, Field, model_validator
from langchain_deepseek import ChatDeepSeek

load_dotenv(verbose=True)

api_key = os.getenv("DEEPSEEK_API_KEY")
api_base = os.getenv("DEEPSEEK_API_BASE")


llm = ChatDeepSeek(
    temperature=0.8,
    model="deepseek-chat",
    api_key=api_key,
    api_base=api_base
    )


class Joke(BaseModel):
    setup: str = Field(description="笑话中的铺垫问题，必须以？结尾")
    punchline: str = Field(description="笑话中回答铺垫问题的部分，通常是一种抖包袱方式回答，如谐音梗、会错意")


parser = PydanticOutputParser(pydantic_object=Joke)

prompt = PromptTemplate(
    template="生成一个笑话，包含铺垫问题和回答部分。\n{format_instructions}\n{text}",
    input_variables=["text"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)

# print(prompt.format(text="请给我讲一个笑话"))

chain = prompt | llm | parser

result = chain.invoke("请给我讲一个笑话")

print(result)